package fem2.material;

import inf.jlinalg.Array1DMatrix;
import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;
import inf.math.LinAlg;

import java.util.ArrayList;
import java.util.Map;

import math2.FunctionRnToR;
import math2.LinAlgUtilities;
import fem2.Element;
import fem2.MaterialModel;
import fem2.Model;
import fem2.Node;
import fem2.NotImplementedError;
import fem2.Patch;
import fem2.element.StructuralElement;
import fem2.enu.StressType;
import fem2.mesh_and_topology.IntegrationPointIndexer;
import fem2.mesh_and_topology.SpatialCriteria;
import fem2.mesh_and_topology.SphereSpatialCriteria;

/**
 * host material model for stress recovery using Moving Least Square weighing
 * function and discrete approach
 * 
 * @author hbui
 * 
 */
public class MLSrecoveryMM extends StressRecoveryMM {

	private IntegrationPointIndexer indexer;
	private double radii;
	private FunctionRnToR weighingFunction;

	/**
	 * stress recovery using moving least square weighing function
	 * 
	 * @param m
	 * @param mm
	 * @param radii
	 */
	public MLSrecoveryMM(Model m, MaterialModel mm, double radii) {
		this(new IntegrationPointIndexer(m.getMesh()), mm, radii);
	}

	public MLSrecoveryMM(IntegrationPointIndexer indexer, MaterialModel mm, double radii) {
		super(mm);
		this.indexer = indexer;
		this.radii = radii;

		weighingFunction = new FunctionRnToR() {

			@Override
			public double valueAt(double... x) {
				double s = x[0];
				if (Math.abs(s) > 1) {
					return 0;
				} else {
					return 1 - 6 * Math.pow(s, 2) + 8 * Math.pow(s, 3) - 3 * Math.pow(s, 4);
				}
			}

			@Override
			public double[] gradientAt(double... x) {
				/*
				 * TODO
				 */
				throw new NotImplementedError();
			}
		};
	}

	/**
	 * stress recovery using specified weighing function
	 * 
	 * @param m
	 * @param mm
	 * @param radii
	 * @param weighingFunction
	 */
	public MLSrecoveryMM(Model m, MaterialModel mm, double radii, FunctionRnToR weighingFunction) {
		super(mm);
		indexer = new IntegrationPointIndexer(m.getMesh());
		this.radii = radii;
		this.weighingFunction = weighingFunction;
	}

	@Override
	public IVector getStress(Element e, int k, StressType stressType) {
		switch (stressType) {
		case RECOVER_STRESS_MLS:
			return computeRecoverStress(e, e.getIntegrationPoint(k));
		default:
			return getMaterialModel().getStress(e, k, stressType);
		}
	}

	@Override
	public void Initialize(Element e) {
		getMaterialModel().Initialize(e);
		indexer.index(e);
	}

	/**
	 * compute recover stress at node
	 * 
	 * @param n
	 * @return
	 */
	public IVector computeRecoverStress(Node n) {

		int dim = n.getDimension();
		int ntrs = dim * (dim + 1) / 2;
		double[] pos = n.getPosition();
		SpatialCriteria sc;

		if (dim == 2) {
			sc = new SphereSpatialCriteria(pos[0], pos[1], radii);
		} else if (dim == 3) {
			sc = new SphereSpatialCriteria(pos[0], pos[1], pos[2], radii);
		} else {
			throw new Error("invalid dimension");
		}

		/*
		 * collect all the integration points around the node
		 */
		Map<Element, ArrayList<Integer>> ips = indexer.queryRange(sc);

		/*
		 * count number of integration points to select appropriate polynomial
		 * order
		 */
		int ng = 0;
		for (Element e : ips.keySet()) {
			ng += ips.get(e).size();
		}
		/*
		 * TODO: check if 6 is the minimal number of points to ensure the sum
		 * matrix is not singular
		 */
		int p = 1;
		if (ng >= 6) {
			p = 2;
		}
		if (ng < 3) {
			throw new Error(
					"not enough integration points to compute left hand side matrix. The result may not be accurate");
		}

		/*
		 * iterate over all integration points to compute left hand side matrix
		 * and right hand side vector
		 */
		IMatrix M = null;
		IVector[] rhs = new IVector[ntrs];
		for (Element e : ips.keySet()) {
			for (int k : ips.get(e)) {
				double[] xi = e.getIntegrationPoint(k);
				double[] X = e.getMeshPart().XAt(xi);
				IVector sigma = ((StructuralElement) e).getStress(k);

				// Debugger.watch("element id: ", e.getId());
				// Debugger.watch("gp: ", k);
				// Debugger.watch("stress: ", sigma);

				double r = LinAlg.diffNorm2(dim, X, pos);
				double w = weighingFunction.valueAt(r / radii);

				IVector P = Patch.computeInterpolationVector(p, false, X);

				if (M == null) {
					M = Array1DMatrix.createRowMajor(P.getSize(), P.getSize());
				}
				BLAM.multiply(w, P, P, M);

				for (int i = 0; i < ntrs; i++) {
					if (rhs[i] == null) {
						rhs[i] = new ArrayVector(P.getSize());
					}
					BLAM.add(w * sigma.get(i), P, rhs[i]);
				}
			}
		}

		// Debugger.watch("M = ", M);
		// for (int i = 0; i < ntrs; i++) {
		// Debugger.watch("rhs " + i + " = ", rhs[i]);
		// }

		/*
		 * solve for each array of coefficients and compute stress component at
		 * node
		 */
		IVector rStress = new ArrayVector(ntrs);
		IVector P = Patch.computeInterpolationVector(p, false, pos);
		for (int i = 0; i < ntrs; i++) {
			LinAlgUtilities
					.backSolve((Array1DMatrix) M.clone(), ((ArrayVector) rhs[i]).getValues());
			rStress.set(i, BLAM.dot(P, rhs[i]));
		}

		return rStress;
	}
}
